# 本脚本的目的：
#   对同一目标 (targetName) 在不同滤光片 (filterName) 下拍摄的多幅单次曝光图像进行叠加处理。
#
# 流程步骤：
#   (1) 基础准备：读取配置文件、图像列表，确认所需目录存在等。
#   (2) 图像测光：对每幅单次曝光图像执行 SExtractor，提取星表以获取星等、FWHM 等信息。
#   (3) 相对光度校准：从所有图像中选取 FWHM 最小的一幅作为参考图像，对其他图像进行归一化。
#   (4) 叠加：使用 Swarp 将经过校准的多幅图像叠加为一张最终图像。
#   (5) 极限星等计算：对叠加后的图像再次测光，得到极限星等，并将结果写入文件。
#
# 本示例还演示了如何在同一个 'with' 语句中打开并写入两个文件：一个记录测光结果 (target_name)，
# 另一个记录光度归一化时的零点信息 (target_name_zp)。

import shutil
from loguru import logger
import numpy as np
from astropy.stats import sigma_clip
from astropy.io import fits
import matplotlib.pyplot as plt
from tools import PSFPhotometry, utils
import os
from tools.objectlist import crdList
from tools.util_statck import calculate_limiting_magnitude, calculate_mid_exposure_time, crossmatch, todegree
from datetime import datetime

from tools.utltil_mosic import get_log_file_path_ospath

if __name__ == "__main__":
    # -------------------------------
    # 1. 基础配置
    # -------------------------------
    targetName = "sb24111301" 
    
    # 目标名称，用于后续路径和文件命名

    limitmaglist = []
    # 用于存储在不同滤光片下计算到的极限星等，便于在处理结束后统一查看或统计

    # 定义两个需要写入的结果文件：
    #   1) target_name：主要保存最终的测光结果，如极限星等、曝光时间等
    #   2) target_name_zp：主要保存光度归一化时参考图像、相对光度比等零点信息
    maindir = "."
    
    reuslt_dir="result/" + targetName 
    utils.ensure_directories(reuslt_dir)
    target_name = os.path.join(reuslt_dir,targetName + ".txt")
    target_name_zp = os.path.join(reuslt_dir,targetName + "_zpinfo.txt")
    target_name_mosic = os.path.join(reuslt_dir,targetName + "_mosic.txt")
    imgdir = os.path.join(maindir, "images", targetName)
    products_dir = os.path.join(maindir, "products", targetName)
    confdir = os.path.join(maindir, "config")
    figdir = os.path.join(maindir, "figures",targetName)
    # 若目标目录不存在，自动创建
    utils.ensure_directories(products_dir)
    utils.ensure_directories(figdir)

            
    date_list = {}  
    # 用于记录每个滤光片对应的图像拍摄日期，可能在后续用于校验观测时段等

    # 使用 with 语句在“追加模式”下 (a) 同时打开两个文件
    with open(target_name, 'a') as file, open(target_name_zp, 'a', encoding='utf-8') as zp_file, open(target_name_mosic, 'a', encoding='utf-8') as mosic_file:
        # 在结果文件 (target_name) 中先写入表头信息
        file.write("Start_Time(UT),Mid-Time(UT),Band,Exposure(s),Lim-mag (AB)\n")
        
        # 依次处理 6 个滤光片：u, v, g, r, i, z
        for filterName in ['z',"v", 'r']: #['u',"v", 'g',"r", 'i',"z"]:
            # 日志文件夹，存放日志文件
            log_dir =os.path.join( "logs",targetName)
            os.makedirs(log_dir, exist_ok=True)  # 如果 logs 文件夹不存在则创建
            
            # 使用 Loguru 配置日志文件，指定大小轮转、日志级别等
            
            log_file_path = get_log_file_path_ospath(log_dir, target_name, filterName,type="stack")
            logger.add(log_file_path, rotation="1 MB", level="DEBUG", backtrace=True, diagnose=True)
        

            logger.info(f"##################### stating {targetName} and {filterName}#####################")
            crdTar = crdList[targetName]
            # 根据滤光片类型，设置不同的像素分辨率 (ipixScale)
            if filterName in ['u', 'v', 'g', 'r']:
                ipixScale = 0.429
                crdTar[2]=6160
                crdTar[3]=6144
            elif filterName in ['i', 'z']:
                ipixScale = 0.286
                crdTar[2]=9232
                crdTar[3]=9216
            
            # 从自定义的对象列表中获取该目标的坐标信息和最后叠加图像的尺寸
            
            # SExtractor 的配置与参数文件
            sexParam = os.path.join(confdir, "default.param") 
            sexConf = os.path.join(confdir, "default_limit.sex") #该地方使用的ORI参数，该参数中的DEBLEND_MINCONT设置的较大，不至于识别到许多特别小的源

            # SExtractor 的命令行模板
            sexComd1 = "sex %s -c %s -PARAMETERS_NAME %s -CATALOG_NAME %s "
            sexComd2 = "-BACK_SIZE 64  -DETECT_MINAREA %s -DETECT_THRESH %s -ANALYSIS_THRESH %s "
            sexComdS1 = sexComd1 + sexComd2

            # 读取当前滤光片下的图像列表文件：images_{filterName}.list
            imgListName = os.path.join(imgdir, f"images_{filterName}.list")
            imgList = open(imgListName, "r").read().splitlines()
            nimg = len(imgList)
            logger.info(f"^_^ Total {nimg} {filterName}-band images")

            newImgListName = os.path.join(products_dir, f"imagesDir_{filterName}.list")
            f_newImgList=open(newImgListName, "w")
            # fwhmList：记录每幅图像的 FWHM 用于后续选择参考图像
            fwhmList = {}
            
            # 用于记录日期、曝光时间和中点时间
            date_ = []
            exposure_value_ = []
            mid_times = []
            
            # 读取每幅图像头中的曝光时间(EXPOSURE)和拍摄时间(DATE)，并计算中点时间
            for i in range(nimg):
                iimgX = imgList[i]
                iimg = os.path.join(imgdir, iimgX)
                
                #将数据拷贝到产品目录
                
                shutil.copy(iimg, products_dir)
                f_newImgList.write(os.path.join(products_dir,iimgX)+"\n")
                
                with fits.open(iimg) as hdul:
                    header = hdul[0].header
                    exposure_value = header.get('EXPOSURE')
                    date_value = header.get('DATE')
                    
                exposure_value_.append(exposure_value)
                date_.append(date_value)
                # calculate_mid_exposure_time：自定义函数，根据开始时间和曝光量算出中点时刻
                mid_exposure_time = calculate_mid_exposure_time(date_value, exposure_value)
                mid_times.append(mid_exposure_time)
            
            # 将该滤光片的日期列表存入字典
            date_list[filterName] = date_
            f_newImgList.flush()
            
            # 对每幅图像调用 SExtractor，提取 FWHM
            for i in range(nimg):
                iimgX = imgList[i]
                iimg = os.path.join(imgdir, iimgX)
                icatX = iimgX[:-4] + "ldac"
                icat = os.path.join(products_dir, icatX)
                
                # 拼装命令并执行
                isex = sexComdS1 % (iimg, sexConf, sexParam, icat,5,1.5,1.5)
                logger.info(isex)
                os.system(isex)
                
                # 从 ldac 文件中读取测光结果 (ext=2)
                ildac = fits.getdata(icat, ext=2)
                logger.info(f"icat:{icat}")
                iflag = ildac["FLAGS"]
                isnr = ildac["SNR_WIN"]
                ifwhm = ildac["FWHM_IMAGE"]
                
                # 筛选：FLAGS=0 且 SNR>20 的源，再取这些源的 FWHM 中位数
                gid = (iflag == 0) & (isnr > 20)
                ifwhm = ifwhm[gid]
                ifwhm = np.median(ifwhm)
                fwhmList[iimgX] = ifwhm
            
            logger.info(f"fwhmList:{fwhmList}")
            # -------------------------------
            # 2. 相对光度校准
            # -------------------------------
            logger.info("^_^ Perform relative photometric calibration")
            
            # 选取 FWHM 最小的图像作为参考图像
            refImg = min(zip(fwhmList.values(), fwhmList.keys()))[1]
            
            # 将选取的参考图像名写入 zp_file 文件中，记录该信息
            zp_file.write(f"choose refImg is {refImg} \n")
            
            refLdacn = os.path.join(products_dir, refImg[:-4] + "ldac")
            refLdac = fits.getdata(refLdacn, ext=2)
            
            # 对参考图像进行简单筛选，用于做光度标定的参考星
            flagRef = refLdac["FLAGS"]
            snrRef = refLdac["SNR_WIN"]
            fluxRef = refLdac["FLUX_AUTO"]
            gid = (flagRef == 0) & (snrRef > 10) & (fluxRef > 0.0)
            fluxRef = fluxRef[gid]
            raRef = refLdac["ALPHA_J2000"][gid]
            decRef = refLdac["DELTA_J2000"][gid]
            nstarRef = len(fluxRef)
            
            logger.info(f"    Selected reference image is {refImg}")
            logger.info(f"    Total {nstarRef} stars are selected for calibration\n")
            zp_file.write(f"Total {nstarRef} stars are selected for calibration\n")
            
            # 读取参考图像的头部信息，用于读取零点等信息
            refImg_abs = os.path.join(imgdir, refImg)
            refimgMat, refhdr = fits.getdata(refImg_abs, header=True)
            
            # 通过目标坐标，确定参考图像上对应的 Gate 分区，继而从头部读出对应的零点 (ZPMAG_G)
            (x, y), (height, width) = utils.ra_dec_to_xy_and_size(refImg_abs, crdTar[0], crdTar[1])
            partitioner = utils.MatrixPartitioner(height, width)
            logger.info(f"MatrixPartitioner:{partitioner}")
            range_info = partitioner.find_submatrix_with_range(x, y)
            
            # 参考图像头部可能存储了多个 Gate 的零点 (ZPMAG_G)，这里根据分区索引取对应零点
            v_list = utils.read_header_keys(refImg_abs, ["exposure","ZPMAG","DATE","ZPMAG_G"])
            gate_find = range_info['submatrix_index'] - 1
            mag_zp = float(v_list['ZPMAG_G'].split(',')[gate_find])
            
            logger.info(f"mag_zp:{mag_zp},v_list['ZPMAG_G']:{v_list['ZPMAG_G']},gate_find:{gate_find},(x, y):{(x, y)}")
            zp_file.write(f"mag_zp:{mag_zp},v_list['ZPMAG_G']:{v_list['ZPMAG_G']},ZPMAG:{v_list['ZPMAG']},gate_find:{gate_find},(x, y):{(x, y)}\n")
            
            
            # 对其他图像进行相对光度归一化
            for iimgKey in fwhmList.keys():
                if iimgKey == refImg:
                    # 参考图像本身，不需要进行光度比计算，可默认为 1
                    zp_file.write(f"{iimgKey},{1}\n")
                    continue
                
                logger.info(f"    Calibrate image {iimgKey}")
                ildacn = os.path.join(products_dir, iimgKey[:-4] + "ldac")
                ildac = fits.getdata(ildacn, ext=2)
                
                iflag = ildac["FLAGS"]
                isnr = ildac["SNR_WIN"]
                iflux = ildac["FLUX_AUTO"]
                gid = (iflag == 0) & (isnr > 20) & (iflux > 0.0)
                iflux = iflux[gid]
                ira = ildac["ALPHA_J2000"][gid]
                idec = ildac["DELTA_J2000"][gid]
                instar = len(iflux)
                logger.info(f"    Total {instar} stars are selected")

                # crossmatch：自定义函数，将参考图像星表与当前图像星表进行交叉匹配
                idRef, iid = crossmatch(raRef, decRef, ira, idec, aperture=2.0)
                
                # 计算相对光度比：参考图像 fluxRef / 当前图像 flux
                irflux = fluxRef[idRef] / iflux[iid]
                irfluxNew = sigma_clip(irflux, sigma=3, maxiters=5, masked=False)
                imrflux = np.nanmedian(irfluxNew)
                istd = np.nanstd(irfluxNew)
                logger.info(f"    The estimated relative flux ratio is {imrflux:6.4f}")

                # 绘制光度比与星等的关系图，并保存到 figures 目录
                magLim = [-2.5 * np.log10(iflux.max()) - 0.5, 
                        -2.5 * np.log10(iflux.min()) + 0.5]
                figName = os.path.join(figdir, "photScale" + iimgKey[:-4] + "png")
                plt.plot(magLim, [imrflux, imrflux], "r-", linewidth=1.5,
                        label="$s=%.3f \pm %.3f$" % (imrflux, istd))
                plt.plot(magLim, [imrflux + istd, imrflux + istd], "r--", linewidth=1.5)
                plt.plot(magLim, [imrflux - istd, imrflux - istd], "r--", linewidth=1.5)
                plt.plot(-2.5 * np.log10(fluxRef[idRef]), irflux, "k.", ms=8)
                plt.xlabel("AUTO MAG", fontsize=10)
                plt.ylabel("FLUX SCALE", fontsize=10)
                plt.legend(loc="upper left", fontsize=12)
                plt.xlim(magLim)
                plt.ylim([imrflux - 5.0 * istd, imrflux + 5.0 * istd])
                plt.savefig(figName, dpi=300)
                plt.clf()
                plt.close()

                # 将光度比例 (imrflux) 写入图像头部，便于后续 Swarp 在叠加时进行归一化
                logger.info("    Update the image header")
                iimg = os.path.join(imgdir, iimgKey)
                iimgMat, ihdr = fits.getdata(iimg, header=True)
                
                #######################保存到另外一个目录
                products_iimg = os.path.join(products_dir, iimgKey)
                ihdr["FLXSCALE"] = imrflux
                
                fits.writeto(products_iimg, iimgMat, ihdr, overwrite=True)
                
                # 同时在 zp_file 中记录此图像的光度比
                zp_file.write(f"{products_iimg},{imrflux}\n")
            
            zp_file.write("\n\n")
            # 每处理完一个滤光片，就将 zp_file 中的内容刷新到磁盘
            zp_file.flush()
            

            # -------------------------------
            # 3. 利用 Swarp 进行叠加
            # -------------------------------
            swarpConf = os.path.join(confdir, "default.swarp")
            swarpComd1 = "swarp @%s -c %s -IMAGEOUT_NAME %s -WEIGHTOUT_NAME %s "
            swarpComd2 = f"-RESAMPLE_DIR %s -PIXEL_SCALE {ipixScale} -SUBTRACT_BACK N -VERBOSE_TYPE FULL "
            swarpComd3 = f"-CENTER_TYPE MANUAL -CENTER {crdTar[0]},{crdTar[1]} -IMAGE_SIZE {crdTar[2]},{crdTar[3]} "
            swarpComd = swarpComd1 + swarpComd2 + swarpComd3
            
            logger.info("^_^ Stack the images ...")
            mosaicName = os.path.join(products_dir, f"zMosaic_{targetName}_{filterName}.fits")
            mosaicWghName = mosaicName[:-4] + "wgh.fits"
            
            # 组装 swarp 命令并执行
            swp = swarpComd % (newImgListName, swarpConf, mosaicName, mosaicWghName, products_dir)
            logger.info(swp)
            os.system(swp)
            
            mosic_file.write(mosaicName+"\n")
            mosic_file.flush()
            
            # 删除 swarp 生成的权重图
            os.system(f"rm {mosaicWghName}")
            

            # -------------------------------
            # 4. 对叠加后的图像做测光，计算 FWHM
            # -------------------------------
            mosaicCat = os.path.join(products_dir, f"zMosaic_{targetName}_{filterName}.ldac")
            mosaicFits = os.path.join(products_dir, f"zMosaic_{targetName}_{filterName}.fits")
            
            # 再次用 SExtractor 提取星表
            isex = sexComdS1 % (mosaicName, sexConf, sexParam, mosaicCat,5,0.5,1)
            os.system(isex)
            
            ildac = fits.getdata(mosaicCat, ext=2)
            iflag = ildac["FLAGS"]
            isnr = ildac["SNR_WIN"]
            ifwhm = ildac["FWHM_IMAGE"]
            
            gid = (iflag == 0) & (isnr > 20)
            ifwhm = ifwhm[gid]
            ifwhm = np.median(ifwhm)
            
            # 将叠加后的 FWHM 写入头部
            iimgMat, ihdr = fits.getdata(mosaicName, header=True)
            ihdr["FWHM"] = (ifwhm, "FWHM of stacked image [pixel]") 
            fits.writeto(mosaicName, iimgMat, ihdr, overwrite=True)
            logger.info(f'the FWHM of stackimg in {filterName} band is {ifwhm}')
            
            # -------------------------------
            # 5. 计算极限星等并写入结果
            # -------------------------------
            # 转换目标坐标字符串为度
            ra_dec = todegree(crdTar[0], crdTar[1])
            
            # 自定义函数 calculate_limiting_magnitude 用于计算叠加图的极限星等
            limiting_mag_info = calculate_limiting_magnitude(
                mosaicName, mosaicCat, ra_dec, filterName,
                mag_zp, 10, mosaicFits, filterName,ipixScale
            )
            logger.info(limiting_mag_info)
            limitmaglist.append(limiting_mag_info)
            
            # 分析所有图像的最早拍摄开始时间 (Start_Time) 和中点时间
            start_times = [datetime.fromisoformat(time) for time in date_]
            min_ut = min(start_times)
            min_ut_str = min_ut.strftime('%Y-%m-%dT%H:%M:%S')
            
            # 取最早和最晚的中点时间做一个平均，得到本滤光片整合后的观测中点时间
            first_time = min(mid_times)
            last_time_exposure = max(mid_times)
            midpoint_time = (
                first_time + (last_time_exposure - first_time) / 2
            ).strftime('%Y-%m-%dT%H:%M:%S')
            
            # limiting_mag_info：
            #   [0] 滤光片名称
            #   [1] 滤光片
            #   [2] 曝光时间 (单张)
            #   [3] 叠加后零点
            #   [4] 其它信息
            #   [5] 计算得到的极限星等
            #
            # 这里将曝光时间写成 limiting_mag_info[2] * nimg，表示将多张图像的总曝光时间累加
            static_single = (
                f"{min_ut_str},{midpoint_time},"
                f"{limiting_mag_info[1]},{limiting_mag_info[2]}*{nimg},"
                f"{limiting_mag_info[5]}\n"
            )
            
            # 写入结果文件
            file.write(static_single)
            logger.info(static_single)
            file.flush()
            
            if True:
                #进一步进行测光和生成reg
                utils.CatalogueRegion(mosaicCat)
                utils.save_region(crdTar[0], crdTar[1],name="targetName", filename=os.path.join('result', targetName, f"{targetName}.reg"))
                
                catalog_data = fits.getdata(mosaicCat, ext=2)
                matched_star=utils.TargetXCatalogue(ra_dec[0], ra_dec[1],catalog_data,radius=2)
                if matched_star: #如果给定的位置是精确的话，但是一般应该都不够精准
                    
                                
                    # convert_to_unit_time_magnitude 用于将测得的 MAG_AUTO 转换成单位时间星等
                    m_per_sec, sigma_m_per_sec = utils.convert_to_unit_time_magnitude(
                        matched_star['MAG_AUTO'],
                        v_list["exposure"],
                        sigma_m=matched_star['MAGERR_AUTO']
                    )
                    # 进一步校正，将 Gate 零点加到单位时间星等上，得到最终校正星等
                    CALMAG = mag_zp + m_per_sec
                    
                    # 计算单位时间星等
                    flux_ = matched_star["FLUX_AUTO"]
                    flux_S_PAN = -2.5 * np.log10(flux_ / v_list["exposure"])
                    logger.info(f"mag_zp:{mag_zp},m_per_sec: {m_per_sec:.4f},CALMAG:{CALMAG},flux_s:{flux_S_PAN}  星等误差: {sigma_m_per_sec:.4f}")
                    
                    
                else:
                    logger.error(f"no star at this position({crdTar[0]},{crdTar[1]})")
                    
                    
                # # 可以在实例化时自定义探测阈值
                # psf_photometry = PSFPhotometry.PSFPhotometry(
                #     mosaicName,
                #     maindir,
                #     detect_minarea=3,      # 自定义 DETECT_MINAREA
                #     detect_thresh=0.5,      # 自定义 DETECT_THRESH
                #     analysis_thresh=0.5     # 自定义 ANALYSIS_THRESH
                # )
                # psf_photometry.execute()
                # logger.info(psf_photometry.ldac_file)
                # catalog_data_psf = fits.getdata(psf_photometry.ldac_file, ext=2)

                # utils.CatalogueRegion(psf_photometry.ldac_file)
                # matched_star=utils.TargetXCatalogue(ra_dec[0], ra_dec[1],catalog_data_psf)
                
                # if matched_star: #如果给定的位置是精确的话，但是一般应该都不够精准
                #     # logger.info(matched_star)
                #     # logger.info(f"statck_mag:{mag_zp+matched_star['MAG_PSF']}")
                                
                #     # convert_to_unit_time_magnitude 用于将测得的 MAG_AUTO 转换成单位时间星等
                #     m_per_sec, sigma_m_per_sec = utils.convert_to_unit_time_magnitude(
                #         matched_star['MAG_PSF'],
                #         v_list["exposure"],
                #         sigma_m=matched_star['MAGERR_PSF']
                #     )
                #     # 进一步校正，将 Gate 零点加到单位时间星等上，得到最终校正星等
                #     CALMAG = mag_zp + m_per_sec
                    
                #     # 计算单位时间星等
                #     flux_ = matched_star["FLUX_PSF"]
                #     mag_S_PAN = -2.5 * np.log10(flux_ / v_list["exposure"])
                #     logger.info(f"mag_zp:{mag_zp},m_per_sec: {m_per_sec:.4f},PSF_CALMAG:{CALMAG},flux_s:{mag_S_PAN}  星等误差: {sigma_m_per_sec:.4f}")
                # else:
                #     logger.error(f"no star at this position({crdTar[0]},{crdTar[1]})")

                        

            
    
